import copy
import numpy as np

from GA import GA
## 继承SA 重新下面两个方法

class myGA(GA):
    def __init__(self, N=4, precison=5, mutation_rate=0.1,var_num=3) -> None:
        super().__init__(N=N, precison=precison, mutation_rate=mutation_rate)
        self.low=None
        self.high=None
        self.best_value_list=[]
        self.var_num=3
    def _encoding_this(self,samples):
        res=[]
        for items in samples:
            one=[]
            for item in items:
                one.append(self._encoding(item))

            res.append(one)
        return res

    def _crossover_this(self,samples):
        res=[[],[]]
        for i in range(self.var_num):
            one=self._crossover(samples[0][i],samples[1][i])

            res[0].append(one[0])
            res[1].append(one[1])

        return res
    
    def _mutation_this(self,samples):
        res=[[],[]]
        for i in range(self.var_num):
            one=self._mutation([samples[0][i],samples[1][i]])

            res[0].append(one[0])
            res[1].append(one[1])

        return res

    def _decoding_this(self,samples):
        res=[]
        for items in samples:
            one=[]
            for item in items:
                one.append(self._decoding(item))

            res.append(one)
        return np.array(res)


    def _update_population_this(self,samples):
        values=[self.get_value(x) for x in samples]
        result=[]
        for i,items in enumerate(self.population):
            dict_data=copy.deepcopy(items)
            dict_data[samples[0][i]]=values[0]
            dict_data[samples[1][i]]=values[1]

            res=sorted(dict_data.items(),key=lambda x:x[1])[:self.N]

            res=dict(res)
            result.append(res)

        return result


    
    def _decoding(self, bin_x):
        x=self.low+(self.high-self.low)*(int(bin_x,2)/(2**self.precision-1))
        return x

    def get_value(self, x):
        x=np.array(x,dtype=np.float32)
        # x=float(x)
        x1=x[0]
        x2=x[1]-x[2]
        y=-(2*x1**2+2*x1+x2**2+3*x2+2)

        return y

    def _seed_children_flag(self,samples):
        for items in samples:
            x=[]
            for item in items:
                one=self._decoding(item)
                if one<0:
                    return False
                x.append(one)
            x1=x[0]
            x2=x[1]-x[2]
            f1=2*x1**2+x1+x2**2+x2
            f2=x1**2+x1+2*x2**2+x2
            f3=x1+2*x2
            f4=x1+x2
            if f1>52:
                return False
            elif f2>8:
                return False
            elif f3<1:
                return False
            elif f4>19:
                return False
            # x3=x[3]-x[4]
            # f1=x1+2*x1**2+x2+2*x2**2+x3
            # f2=x1+x1**2+x2+x2**2-x3
            # f3=2*x1+x1**2+2*x2+x3
            # f4=x1**2+x3
            # f5=x1+2*x2
            # if f1>10:
            #     return False
            # elif f2>50:
            #     return False
            # elif f3>40:
            #     return False
            # elif f4<2-1 or f4>2+1:
            #     return False
            # elif f5<1:
            #     return False

        return True

    def _initiate_population(self, xs):
        res=[]
        xs_new=copy.deepcopy(xs)
        for _ in range(self.N):
            x=np.random.choice(xs_new)
            res.append(x)
            xs_new.remove(x)

        return res


    def initate_value(self,population):

        l=len(population)
        populations=[]
        for i in range(l):
            populations.append({})
        
        for i in range(self.N): 
            one=[]
            for it in population:
                one.append(it[i])

            value=self.get_value(one)
            for j,it in enumerate(population):
                populations[j][it[i]]=value

        return populations




    def create_x(self,population):
        res=[]
        for i in range(self.N):
            one=[]
            for j in range(self.var_num):
                one.append(population[j][i])

            res.append(one)

        return np.array(res)

    def get_samples(self, population):
        res=[]
        for i in range(2):
            one=[]
            for j,it in enumerate(population):
                its=list(it.keys())
                
                if i==0:
                    one.append(np.random.choice(its))
                else:
                    one.append(np.random.choice(list(set(its)-set({res[0][j]}))))

            res.append(one)

        return np.array(res)

            

    def __call__(self,low,heigh, epoch=100):
        self.low=low
        self.high=heigh
        self.flag=0
        self.population=[]
        self.best_x=None
        self.traces=[]
        self.all_population=[i for i in range(low,int(heigh)+1)]
        popus=[]
        for i in range(self.var_num):
            popus.append(self._initiate_population(self.all_population))

        self.population=self.initate_value(popus)
        self.best_value=1000
        
        for _ in range(epoch):
            
            samples=self.get_samples(self.population)

            # bin_x=[self._encoding(res) for res in samples]
            bin_x=self._encoding_this(samples)
            while True:
                x_row=self._crossover_this(bin_x)
                x=self._mutation_this(x_row)

                flag=self._seed_children_flag(x)

                if flag:
                    break
            
            # xs=[self._decoding(res) for res in x]
            xs=self._decoding_this(x)
            # one=[]
            for i in range(2):
                value=self.get_value(xs[i])
                # one.append(np.append(xs[i],value))
                
                # print(value,end="\t")
                if  value<self.best_value:
                    self.best_value=copy.deepcopy(value)
                    self.best_x=copy.deepcopy(xs[i])
            # print()
            self.traces.append(copy.deepcopy(self.population))
            self.population=self._update_population_this(xs)
            print(self.population)
            
            
            # self.best_value_list.append(copy.deepcopy(self.population))
            # self.best_value=min(list(self.population.values()))


        print(self.best_value)
            




if __name__=="__main__":
    from plt_helper import plot_dynamic_figure

    # data={"T1":'1235','T2':'2145','T3':'1543','T4':'4253','T5':'2531'}

    
    # model=My_SA()
    # model(data=data,x0='12345',epoch=5)
    model=myGA(precison=6)
    model(0,6,epoch=100)

    traces=model.traces
    res=[]

    for its in traces:
        one=[]
        one_deal=[]
        for it in its:
            keys=list(it.keys())
            values=list(it.values())
            # for item in values:
            #     keys.append(item)
            one.append(keys)
        one=np.array(one)

        for i in range(model.N):
            one_item=list(one[:,i])
            one_item.append(values[i])

            one_deal.append(one_item)

        res.append(one_deal)

    res=np.array(res)

    print(res.shape)



    def fun(x1,x2):
        value=2*x1**2+2*x1+x2**2+3*x2+2

        return value

    plot_dynamic_figure(res,fun)

    # x_source=np.linspace(0,2*np.pi,320)
    # y_source=-np.sin(x_source)

    

    # trace=model.best_value_list
    # fig,ax=plt.subplots()
    
    # plt.axis([0,31,0,-1000])
    # plt.ion()
    # for item in trace:
    #     xs=[float(x) for x in list(item.keys())]
    #     ys=[float(y) for y in list(item.values())]
    #     ax.cla()
    #     ax.plot(x_source,y_source)
    #     ax.scatter(xs,ys,c='red')
        
    #     plt.pause(0.5)

    # plt.show()



